package chatMemory

import (
	"time"

	"github.com/tmc/langchaingo/llms"
)

var chatMemoryMap map[string]*chatMemoryContent

type chatMemoryContent struct {
	time     int64
	messages []llms.MessageContent
}

func init() {
	chatMemoryMap = make(map[string]*chatMemoryContent)
}

func Put(session string, chatMessage llms.MessageContent) {
	memoryContent, exists := chatMemoryMap[session]
	if !exists {
		memoryContent = &chatMemoryContent{messages: make([]llms.MessageContent, 0)}
	}
	memoryContent.time = time.Now().Unix()
	messages := append(memoryContent.messages, chatMessage)
	memoryContent.messages = messages
	chatMemoryMap[session] = memoryContent
}

func BatchPut(session string, messages ...llms.MessageContent) {
	for _, chatMessage := range messages {
		Put(session, chatMessage)
	}
}

func Messages(session string) []llms.MessageContent {
	memoryContent := chatMemoryMap[session]
	memoryContent.time = time.Now().Unix()
	chatMemoryMap[session] = memoryContent
	return memoryContent.messages
}
